/*
 * factor_graph.h
 *
 *  Created on: Jun 2, 2011
 *      Author: tqlong
 */

#ifndef FACTOR_GRAPH_H_
#define FACTOR_GRAPH_H_

#include <vector>
#include <map>

#include "universe.h"
#include "argument.h"
#include "assignment.h"
#include "factor.h"

/** A set of factors on a set of random variables (universe)
 *  determines a factor graph
 */
template<typename Function>
  class FactorGraph : protected std::vector<Function>
  {
  public:
    typedef Function func_type;
    typedef std::vector<Function> __Base;
    typedef typename func_type::number_type number_type;
    typedef std::vector<int> vint_type;
  protected:
    Argument argument_;
    std::map<int, int> argMap_;
  public:

    /** add a new factor
     *  add new argument if needed and prepare argument mapping
     */
    FactorGraph&
    addFactor(const func_type& f)
    {
      push_back(f);
      Argument f_arg = f.argument();
      for (int v_idx = 0; v_idx < f_arg.n_var(); v_idx++)
        {
          int varIdx = f_arg.varIdx(v_idx);
          std::map<int, int>::iterator it = argMap_.find(varIdx);
          if (it == argMap_.end()) // a new variable
            {
              argument_.addVariable(varIdx);
              argMap_[varIdx] = argument_.n_var() - 1;
            }
        }
      return *this;
    }

    int
    n_factor() const
    {
      return this->size();
    }

    const func_type&
    factor(int i) const
    {
      return this->at(i);
    }

    /** compute product of factors with assignment in the universe */
    number_type
    factorProduct(const Universe& u) const
    {
      number_type product(1.0);
      for (int i = 0; i < n_factor(); i++)
        product *= this->at(i).value(u);
      return product;
    }

    /** compute product of factors with an assignment to variables */
    number_type
    factorProduct(const Assignment& a) const
    {
      number_type product(1.0);
      for (int i = 0; i < n_factor(); i++)
        {
          const func_type& f = this->at(i);
          Assignment f_a = factorAssignment(f, a);
          product *= f.value(f_a);
        }
      return product;
    }

    /** generate an assignment of a factor, make it compatible with the given assignment */
    Assignment
    factorAssignment(const func_type& f, const Assignment& a) const
    {
      const Argument& f_arg = f.argument();
      Assignment f_a = f.firstAssignment();
      for (int v_idx = 0; v_idx < f_arg.n_var(); v_idx++)
        {
          int varIdx = f_arg.varIdx(v_idx);
          std::map<int, int>::const_iterator it = argMap_.find(varIdx);
          if (it == argMap_.end())
            {
              std::cerr << "Variable not found " << __FILE__ << ":" << __LINE__
                  << std::endl;
              exit(1);
            }
          int argPos = it->second;
          f_a[v_idx] = a[argPos];
        }
      return f_a;
    }

    const Argument&
    argument() const
    {
      return argument_;
    }

    Assignment
    firstAssignment() const
    {
      return Assignment(argument_.n_var());
    }

    /** Compute Z, naively */
    number_type
    normalizingConstant(const Universe& u) const
    {
      Argument arg = argument();
      Assignment a = firstAssignment().fixed(arg, u);
      std::cout << "u=" << u << std::endl;
      //std::cout << "arg=" << arg << std::endl;
      number_type Z(0.0);
      for (; a.isValid(); a.nextAssignment(arg, u))
        {
          //std::cout << "a=" << a << std::endl;
          Z += factorProduct(a);
        }
      return Z;
    }

    /** Compute MAP, naively */
    Assignment
    mapInference(const Universe& u) const
    {
      Argument arg = argument();
      Assignment a = firstAssignment().fixed(arg, u);
      std::cout << "u=" << u << std::endl;
      //std::cout << "arg=" << arg << std::endl;
      number_type max(0.0);
      Assignment map;
      for (; a.isValid(); a.nextAssignment(arg, u))
        {
          //std::cout << "a=" << a << std::endl;
          number_type f_prod = factorProduct(a);
          if (max < f_prod)
            {
              max = f_prod;
              map = a;
            }
        }
      return map;
    }
  };

#endif /* FACTOR_GRAPH_H_ */
